#include "rclcpp/rclcpp.hpp"
#include "geometry_msgs/msg/pose_stamped.hpp"
#include "visualization_msgs/msg/marker.hpp"
#include "MyKalmanFilter.h" // 正确的头文件路径
#include <Eigen/Dense>
#include <vector>
#include <boost/any.hpp>

using namespace hitcrt;

class KfInterpolatorNode : public rclcpp::Node
{
public:
    KfInterpolatorNode() : Node("kf_interpolator_node"), is_initialized_(false)
    {
        this->declare_parameter<double>("prediction_rate_hz", 90.0);
        this->declare_parameter<double>("process_noise_pos", 0.05);
        this->declare_parameter<double>("process_noise_vel", 0.5);
        this->declare_parameter<double>("measurement_noise", 0.01);
        this->declare_parameter<std::string>("camera_frame_id", "camera_link");

        prediction_rate_ = this->get_parameter("prediction_rate_hz").as_double();
        double p_n_pos = this->get_parameter("process_noise_pos").as_double();
        double p_n_vel = this->get_parameter("process_noise_vel").as_double();
        measurement_noise_ = this->get_parameter("measurement_noise").as_double();
        camera_frame_id_ = this->get_parameter("camera_frame_id").as_string();

        kf_ = std::make_unique<MyKalmanFilter>();
        
        Eigen::MatrixXd process_noise_cov = Eigen::MatrixXd::Zero(6, 6);
        process_noise_cov.diagonal() << p_n_pos, p_n_pos, p_n_pos, p_n_vel, p_n_vel, p_n_vel;
        
        std::vector<boost::any> init_param = {process_noise_cov};
        kf_->init(init_param);

        subscription_ = this->create_subscription<geometry_msgs::msg::PoseStamped>(
            "/basketball_pose", 10, std::bind(&KfInterpolatorNode::measurement_callback, this, std::placeholders::_1));
        
        filtered_pose_pub_ = this->create_publisher<geometry_msgs::msg::PoseStamped>("/basketball_pose_filtered", 10);
        trajectory_marker_pub_ = this->create_publisher<visualization_msgs::msg::Marker>("/basketball_trajectory_predicted", 10);

        auto timer_period = std::chrono::duration<double>(1.0 / prediction_rate_);
        prediction_timer_ = this->create_wall_timer(timer_period, std::bind(&KfInterpolatorNode::prediction_callback, this));

        RCLCPP_INFO(this->get_logger(), "3D卡尔曼滤波器插帧节点已初始化。");
    }

private:
    void measurement_callback(const geometry_msgs::msg::PoseStamped::SharedPtr msg)
    {
        std::vector<double> measurement_vec = {msg->pose.position.x, msg->pose.position.y, msg->pose.position.z};
        
        Eigen::MatrixXd measurement_noise_cov = Eigen::Matrix3d::Identity() * measurement_noise_;

        rclcpp::Time current_time = this->get_clock()->now();
        double dt = (last_measurement_time_.seconds() > 0) ? (current_time - last_measurement_time_).seconds() : (1.0 / 30.0);
        if (dt <= 0) dt = (1.0/30.0); // 防止dt为0或负数
        last_measurement_time_ = current_time;

        std::vector<boost::any> runtime_params = {dt, measurement_vec, measurement_noise_cov};
        kf_->setParam(runtime_params);

        if (!is_initialized_)
        {
            kf_->setStateByMeasure().resetCovPost();
            is_initialized_ = true;
            RCLCPP_INFO(this->get_logger(), "3D卡尔曼滤波器已使用首次测量值初始化。");
        }
        else
        {
            kf_->predict();
        }
    }

    void prediction_callback()
    {
        if (!is_initialized_) return;

        kf_->stateUpdate();
        
        Eigen::MatrixXd result = kf_->getResult();
        
        auto pose_msg = std::make_unique<geometry_msgs::msg::PoseStamped>();
        pose_msg->header.stamp = this->get_clock()->now();
        pose_msg->header.frame_id = camera_frame_id_;
        pose_msg->pose.position.x = result(0, 0);
        pose_msg->pose.position.y = result(1, 0);
        pose_msg->pose.position.z = result(2, 0);
        pose_msg->pose.orientation.w = 1.0;
        filtered_pose_pub_->publish(std::move(pose_msg));

        publish_trajectory_marker(result.block<3,1>(0,0));
    }

    void publish_trajectory_marker(const Eigen::Vector3d& pos)
    {
        geometry_msgs::msg::Point p;
        p.x = pos.x(); p.y = pos.y(); p.z = pos.z();
        trajectory_points_.push_back(p);

        if (trajectory_points_.size() > 500)
        {
            trajectory_points_.erase(trajectory_points_.begin());
        }

        auto marker_msg = std::make_unique<visualization_msgs::msg::Marker>();
        marker_msg->header.stamp = this->get_clock()->now();
        marker_msg->header.frame_id = camera_frame_id_;
        marker_msg->ns = "predicted_trajectory";
        marker_msg->id = 0;
        marker_msg->type = visualization_msgs::msg::Marker::LINE_STRIP;
        marker_msg->action = visualization_msgs::msg::Marker::ADD;
        marker_msg->points = trajectory_points_;
        marker_msg->scale.x = 0.05;
        marker_msg->color.a = 1.0;
        marker_msg->color.r = 0.0;
        marker_msg->color.g = 1.0;
        marker_msg->color.b = 0.0;
        trajectory_marker_pub_->publish(std::move(marker_msg));
    }

    rclcpp::Subscription<geometry_msgs::msg::PoseStamped>::SharedPtr subscription_;
    rclcpp::Publisher<geometry_msgs::msg::PoseStamped>::SharedPtr filtered_pose_pub_;
    rclcpp::Publisher<visualization_msgs::msg::Marker>::SharedPtr trajectory_marker_pub_;
    rclcpp::TimerBase::SharedPtr prediction_timer_;
    rclcpp::Time last_measurement_time_;

    std::unique_ptr<MyKalmanFilter> kf_;
    bool is_initialized_;
    double prediction_rate_;
    double measurement_noise_;
    std::string camera_frame_id_;
    std::vector<geometry_msgs::msg::Point> trajectory_points_;
};

int main(int argc, char *argv[])
{
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<KfInterpolatorNode>());
    rclcpp::shutdown();
    return 0;
}